from DataContainer.DataModel import DataModel;
from DataContainer.DataView import DataView;
from DataContainer.DataSettings import DataSettings;
from PyQt4 import QtCore, QtGui;
import pickle, os;
import pandas as pd;
import numpy as np;
import pygraphviz as pgv;

class DataController(QtCore.QObject):
    def __init__(self, parent, graphScene):
        super(DataController, self).__init__();
        self.dataView = None;
        self.parent = parent;
        self.mainWin = parent.mainWin;
#        self.graphSettings = parent.graphSettings;
        self.graphScene = graphScene;
        self.dataModel = DataModel();
        self.dataView = DataView(self);
        self.dataSettings = DataSettings(self);
        self.sFilePath = "";
        self.sDir = "./StoredData/";
        self.sExt = "Data Files (*.dat)";
#         self.loadGrpActs_Data(); #TESTING GrpActs Data Table
        
    def setWindowFileTitle(self):
        sFileName = self.sFilePath.split("/");
        sFileName = sFileName[len(sFileName)-1];
        self.mainWin.setWindowTitle(sFileName + " - Visual_B3");
    
    def saveDataAs(self):
        sTitle = "Save Data tables As...";
        self.sFilePath = QtGui.QFileDialog.getSaveFileName(self.mainWin, sTitle, self.sDir, self.sExt);
        if self.sFilePath == "": return;
        self.setWindowFileTitle();
        self.saveData();
    
    def saveData(self):
        sTitle = "Save Data tables...";
        if self.sFilePath == "":
            self.sFilePath = QtGui.QFileDialog.getSaveFileName(self.mainWin, sTitle, self.sDir, self.sExt);
        if self.sFilePath == "": return;
        self.setWindowFileTitle();
        self.saveTables();
    
#     def loadData(self, sFilePath=""):
#         sTitle = "Load Data tables...";
#         if sFilePath!="":
#             self.sFilePath = sFilePath;
#         else:
#             self.sFilePath = QtGui.QFileDialog.getOpenFileName(self.mainWin, sTitle, self.sDir, self.sExt);
#         if self.sFilePath == "": return;
#         self.setWindowFileTitle();
#         self.loadTables();
        
    def loadData(self, sFilePath=""):
        sTitle = "Load Group Activities Data...";
        if sFilePath!="":
            self.sFilePath = sFilePath;
        else:
            self.sFilePath = QtGui.QFileDialog.getExistingDirectory(self.mainWin, sTitle, self.sDir);
        if self.sFilePath == "": return;
        self.sFilePath = str(self.sFilePath.replace("\\", "/"));
        self.loadGrpActs_Data(self.sFilePath);
        
    def loadGrpActs_Data(self, sDirPath="./StoredData/ActGrps_Test"):
        def get_Parents_Roles(sPath):
            fParents = open(sPath+'/parents_roles.txt');
            aParents = fParents.readline().split(',');
            aRoles = fParents.readline().split(',');
            for i,(sP,sR) in enumerate(zip(aParents, aRoles)):
                aParents[i] = int(sP);
                aRoles[i] = int(sR);
            return aParents, aRoles;
        aRoleNames = ['NA', 'FFAER', 'MEETER', 'MEETER-STAND', 'MOVER'];
        aCols = ['Parent-Id', 'Activity', 'Start-Frm', 'End-Frm', 'Role', 'Persons'];
        aParents, aRoles = get_Parents_Roles(sDirPath);
        dfGrpActs = pd.DataFrame(index=range(len(aParents)), columns=aCols);
        for sFile in os.listdir(sDirPath):
            if sFile.startswith('parents_roles'): continue;
            dfActGrp = pd.read_csv(sDirPath+'/'+sFile);
            dfActGrp = dfActGrp.drop('-1',1);
            aGrps = np.unique(dfActGrp.values);
            aGrps = aGrps[aGrps!=-1];
            sAct = sFile.split('.')[:-1][0];
            for iGrp in aGrps:
                if iGrp==aParents[iGrp]: aParents[iGrp] = -1;
                dfGrpActs['Parent-Id'][iGrp] = aParents[iGrp];
                dfGrpActs['Activity'][iGrp] = sAct;
                dfGrpActs['Role'][iGrp] = aRoleNames[aRoles[iGrp]];
                #Get & Set Init, End Frames & Persons involved
                dfGrp = dfActGrp[dfActGrp==iGrp];
                dfGrp = dfGrp.dropna(how='all').dropna(how='all', axis=1);
                dfGrpActs['Start-Frm'][iGrp] = min(dfGrp.index);
                dfGrpActs['End-Frm'][iGrp] = max(dfGrp.index);
                dfGrpActs['Persons'][iGrp] = map(lambda x: int(x), dfGrp.columns);
        self.dataModel.dfGrpActs = dfGrpActs;
        self.dataView.gridGrpActsTbl.updateGrid();
        self.updateGraph2(dfGrpActs);
        
    def saveTables(self):
        aTables = self.dataModel.getDataTables();
        aTables.append(self.getNodesPos());
        pickle.dump(aTables, open(self.sFilePath, "w"));
        
    def loadTables(self):
        aTables = pickle.load(open(self.sFilePath));
        self.dataModel.setDataTables(aTables);
        self.setNodesPos(aTables[len(aTables)-1]);
        self.dataView.updateGridTables();
#        self.dataView.gridStInfo.updateGrid();
#        self.dataView.gridStTrans.updateGrid();
#        self.dataView.gridStDists.updateGrid();
        
    def saveSettings(self):
        self.dataSettings.saveSettings();
        
    def loadSettings(self, sFilePath=""):
        self.dataSettings.loadSettings(sFilePath);
        
    def getNodesPos(self):
        aNodesPos = [];
        for i in range(len(self.graphScene.nodes)):
            iRow = [];
            iRow.append(self.graphScene.nodes[i].pos().x());
            iRow.append(self.graphScene.nodes[i].pos().y());
            aNodesPos.append(iRow);
        return aNodesPos;
    
    def setNodesPos(self, aNodesPos):
        aPosX, aPosY = [], [];
        for i in range(len(aNodesPos)):
            aPosX.append(aNodesPos[i][0]);
            aPosY.append(aNodesPos[i][1]);
        self.graphScene.createNodes(aPosX, aPosY);
    
    def updateDataFromGraph(self):
        edges = self.graphScene.edges;
        self.dataModel.initStTrans();
        if len(edges)>0:
            stTrans = self.graphScene.graphModel.getEdges_TransVals(edges);
            self.dataModel.stateTrans = stTrans;
    
    def updateGraph(self):
        self.graphScene.buildDataGraph(self.dataModel.stateTrans, self.dataModel.stInfo);
        self.updateNodesInfo();
        self.updateSelectedGraphSettings();
        
    def updateGraph2(self, dfGrpActs):        
        aEdges, iCluster = [], 0;
        g = pgv.AGraph(name='ActGrps_Graph', strict=True, directed=True);
        g.graph_attr['newrank']='true';
        g.node_attr['fixedsize']='true';
        #Get PyGraphViz & a List of Edges
        dfGrpActs_str = dfGrpActs.copy();
        dfGrpActs_str['Persons_str'] = dfGrpActs_str['Persons'].apply(str);
        for sP, dfPersons in dfGrpActs_str.groupby(["Persons_str"]):  # @UnusedVariable
            for iP, dfParent in dfPersons.groupby(["Parent-Id"]):
                if iP==-1: continue;
                dfParent.sort("Start-Frm", inplace=True);
                iP_past, aNodes = -1, [];
                for j,iTo in enumerate(dfParent.index):
                    iFrom = iP if j==0 else iP_past;
                    iP_past = iTo;
                    g.add_edge(iFrom, iTo);
                    aEdges.append([iFrom, iTo]);
                    aNodes.append(iTo);
                g.add_subgraph(aNodes, name='cluster_'+str(iCluster), rank='same');
                iCluster += 1;
#         g.draw('./StoredData/Graphs/ActGrps_AGraph.pdf', prog='dot');
        #Get Node Positions
        g.layout(prog='dot');
        aX, aY = [], [];
        for i in range(len(g.nodes())):
            node = g.get_node(i);
            nX, nY = node.attr['pos'].split(",");
            aX.append(float(nX)*3.5-1200);
            aY.append(float(nY)*-3.5+200);
        self.graphScene.graphVizNodes = g;
        self.graphScene.buildDataGraph(aEdges, [aX, aY]);
        self.updateNodesInfo2();
        self.graphScene.showNodeClusters();
        
    def updateNodesInfo2(self):
        nodes = self.graphScene.nodes;
        dfGrpActs = self.dataModel.dfGrpActs;
        self.graphScene.graphModel.setNodes_GrpActsInfo(nodes, dfGrpActs);        
        
    def updateNodesInfo(self):
        aStDims, aRewards, aVals, mQVals = [], [], [], [];
        aPolicy, aStTypes, aActTypes = [], [], [];
        aDists = self.dataModel.selectedDists;
        nDims = self.dataModel.nDims;
        nActs = self.dataModel.nActions;
        nRiskCols = self.dataModel.nRiskActResCols;
        for i in range(len(self.dataModel.stInfo)):
            aStDims.append(self.dataModel.stInfo[i][0:nDims]);
            aRewards.append(self.dataModel.stInfo[i][nDims]);
            aVals.append(self.dataModel.stInfo[i][nDims+1]);
            aPolicy.append(int(float(self.dataModel.stInfo[i][nDims+2])));
            mQVals.append(self.dataModel.stInfo[i][nDims+3:nDims+3+nActs]);
            aStTypes.append(self.dataModel.stInfo[i][nDims+3+nActs]);
        nodes = self.graphScene.nodes;
        iRow = 0;
        for i in range(len(nodes)):
            actTypeRow = [];
            for j in range(len(nodes[i].actions)): #@UnusedVariable
#                print len(self.dataModel.riskActRes), len(self.dataModel.riskActRes[0]), iRow, nRiskCols-1, len(nodes), len(nodes[i].actions);
                actTypeRow.append(self.dataModel.riskActRes[iRow][nRiskCols-1]);
                iRow += 1;
            aActTypes.append(actTypeRow);
        self.graphScene.graphModel.setNodes_StateDim(nodes, aStDims);
        self.graphScene.graphModel.setNodes_Reward(nodes, aRewards);
        self.graphScene.graphModel.setNodes_Value(nodes, aVals);
        self.graphScene.graphModel.setNodes_Distance(nodes, aDists);
        self.graphScene.graphModel.setNodes_StateType(nodes, aStTypes);
        self.graphScene.graphModel.setActions_QVals(nodes, mQVals);
        self.graphScene.graphModel.setActions_ActType(nodes, aActTypes, aPolicy);
        
    def updateSelectedGraphSettings(self):
        self.parent.graphSettings.updateSelectedGraphSettings();
        self.updateBkgColors();
        
    def updateBkgColors(self):
        self.parent.actionCmds.updateBkgColors();
    